/**
 * 
 */
package lv.pig.base.shiro;

import java.util.Arrays;

import javax.servlet.ServletRequest;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import lv.pig.base.jpa.service.UserService;

import org.apache.shiro.codec.Base64;
import org.apache.shiro.subject.PrincipalCollection;
import org.apache.shiro.subject.Subject;
import org.apache.shiro.subject.SubjectContext;
import org.apache.shiro.web.mgt.CookieRememberMeManager;
import org.apache.shiro.web.servlet.Cookie;
import org.apache.shiro.web.servlet.ShiroHttpServletRequest;
import org.apache.shiro.web.servlet.SimpleCookie;
import org.apache.shiro.web.subject.WebSubject;
import org.apache.shiro.web.subject.WebSubjectContext;
import org.apache.shiro.web.util.WebUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

/**
 * @author chenbo
 * 
 */
@Component("myRememberMeManager")
public class MyCookieRememberMeManager extends CookieRememberMeManager
{
    
    private static transient final Logger log = LoggerFactory.getLogger(MyCookieRememberMeManager.class);
    
    public static final String DEFAULT_REMEMBER_ME_COOKIE_PASSWORD = "rememberMe2";
    
    // 保存密码用cookie
    private Cookie enCookie;
    
    @Autowired
     UserService userService;
    
    /**
     * Constructs a new {@code CookieRememberMeManager} with a default {@code rememberMe} cookie template.
     */
    public MyCookieRememberMeManager()
    {
        Cookie enCookie = new SimpleCookie(DEFAULT_REMEMBER_ME_COOKIE_PASSWORD);
        enCookie.setHttpOnly(true);
        enCookie.setMaxAge(Cookie.ONE_YEAR);
        this.enCookie = enCookie;
    }
    
    public Cookie getEnCookie()
    {
        return enCookie;
    }
    
    public void setEnCookie(Cookie enCookie)
    {
        this.enCookie = enCookie;
    }
    
    /**
     * Base64-encodes the specified serialized byte array and sets that base64-encoded String as the cookie value.
     * <p/>
     * The {@code subject} instance is expected to be a {@link WebSubject} instance with an HTTP Request/Response pair
     * so an HTTP cookie can be set on the outgoing response. If it is not a {@code WebSubject} or that
     * {@code WebSubject} does not have an HTTP Request/Response pair, this implementation does nothing.
     *
     * @param subject the Subject for which the identity is being serialized.
     * @param serialized the serialized bytes to be persisted.
     */
    protected void rememberSerializedIdentity(Subject subject, byte[] serialized)
    {
        
        if (!WebUtils.isHttp(subject))
        {
            if (log.isDebugEnabled())
            {
                String msg =
                    "Subject argument is not an HTTP-aware instance.  This is required to obtain a servlet "
                        + "request and response in order to set the rememberMe cookie. Returning immediately and "
                        + "ignoring rememberMe operation.";
                log.debug(msg);
            }
            return;
        }
        
        HttpServletRequest request = WebUtils.getHttpRequest(subject);
        HttpServletResponse response = WebUtils.getHttpResponse(subject);
        
        // base 64 encode it and store as a cookie:
        String base64 = Base64.encodeToString(serialized);
        Cookie template = getCookie(); // the class attribute is really a
                                       // template for the outgoing cookies
        Cookie cookie = new SimpleCookie(template);
        cookie.setValue(base64);
        cookie.saveTo(request, response);
        
        // 增加密码cookie
        String loginName = subject.getPrincipal().toString();
        byte[] password = userService.findByLoginName(loginName).getPassword().getBytes();
        String base64pass = Base64.encodeToString(password);
        Cookie cookiepass = new SimpleCookie(this.getEnCookie());
        cookiepass.setValue(base64pass);
        if (log.isDebugEnabled())
        {
            log.debug("encrypt password:{},password length:{}", base64pass, password.length);
        }
        cookiepass.saveTo(request, response);
    }
    
    public PrincipalCollection getRememberedPrincipals(SubjectContext subjectContext)
    {
        PrincipalCollection principals = null;
        try
        {
            byte[] bytes = getRememberedSerializedIdentity(subjectContext);
            if (bytes != null && bytes.length > 0)
            {
                principals = convertBytesToPrincipals(bytes, subjectContext);
                String loginName = principals.getPrimaryPrincipal().toString();
                if (!checkPassword(loginName, subjectContext))
                {
                    if (log.isDebugEnabled())
                    {
                        log.debug("user:{}'s password in cookie is not matched in database,need Re-login!", loginName);
                    }
                    return null;
                }
            }
        }
        catch (RuntimeException re)
        {
            principals = onRememberedPrincipalFailure(re, subjectContext);
        }
        return principals;
    }
    
    /**
     * check password in cookie
     * 
     * @param loginName
     * @param subjectContext
     * @return
     */
    private boolean checkPassword(String loginName, SubjectContext subjectContext)
    {
        // 校验密码
        WebSubjectContext wsc = (WebSubjectContext)subjectContext;
        if (isIdentityRemoved(wsc))
        {
            return false;
        }
        HttpServletRequest request = WebUtils.getHttpRequest(wsc);
        HttpServletResponse response = WebUtils.getHttpResponse(wsc);
        byte[] password = userService.findByLoginName(loginName).getPassword().getBytes();
        String base64 = getEnCookie().readValue(request, response);
        if (base64 != null)
        {
            base64 = ensurePadding(base64);
            byte[] decoded = Base64.decode(base64);
            if (log.isDebugEnabled())
            {
                log.debug("encrypt password:{}, decode length:{},password length:{}",
                    base64,
                    decoded.length,
                    password.length);
            }
            return Arrays.equals(password, decoded);
        }
        return false;
    }
    
    /**
     * Sometimes a user agent will send the rememberMe cookie value without padding, most likely because {@code =} is a
     * separator in the cookie header.
     * <p/>
     * Contributed by Luis Arias. Thanks Luis!
     *
     * @param base64 the base64 encoded String that may need to be padded
     * @return the base64 String padded if necessary.
     */
    private String ensurePadding(String base64)
    {
        int length = base64.length();
        if (length % 4 != 0)
        {
            StringBuilder sb = new StringBuilder(base64);
            for (int i = 0; i < length % 4; ++i)
            {
                sb.append('=');
            }
            base64 = sb.toString();
        }
        return base64;
    }
    
    private boolean isIdentityRemoved(WebSubjectContext subjectContext)
    {
        ServletRequest request = subjectContext.resolveServletRequest();
        if (request != null)
        {
            Boolean removed = (Boolean)request.getAttribute(ShiroHttpServletRequest.IDENTITY_REMOVED_KEY);
            return removed != null && removed;
        }
        return false;
    }
}
